from pathlib import Path
from PIL import Image
from typing import Tuple, Optional


def ensure_same_height(img1: Image.Image, img2: Image.Image) -> Tuple[Image.Image, Image.Image]:
    """确保两张图片高度一致，若不一致则按比例缩放较矮的图片"""
    if img1.height == img2.height:
        return img1, img2

    # 以较高的图片高度为基准，缩放较矮的图片
    max_height = max(img1.height, img2.height)
    if img1.height < max_height:
        ratio = max_height / img1.height
        new_width = int(img1.width * ratio)
        img1 = img1.resize((new_width, max_height), Image.Resampling.LANCZOS)
    else:
        ratio = max_height / img2.height
        new_width = int(img2.width * ratio)
        img2 = img2.resize((new_width, max_height), Image.Resampling.LANCZOS)
    return img1, img2


def merge_pair(left_img_path: Path, right_img_path: Path, output_path: Path) -> None:
    """将左图（带1）和右图（不带1）合成一张新图，保存到output_path"""
    try:
        # 确保输出目录存在
        output_path.parent.mkdir(parents=True, exist_ok=True)

        # 打开图片
        with Image.open(left_img_path) as left_img, Image.open(right_img_path) as right_img:
            # 确保高度一致
            left_img, right_img = ensure_same_height(left_img, right_img)

            # 创建合成画布（宽度为两图之和，高度为统一后的高度）
            combined_width = left_img.width + right_img.width
            combined_height = left_img.height
            combined_img = Image.new("RGB", (combined_width, combined_height))

            # 粘贴图片（左图居左，右图接在右侧）
            combined_img.paste(left_img, (0, 0))
            combined_img.paste(right_img, (left_img.width, 0))

            # 保存合成图
            combined_img.save(output_path)
            print(f"✅ 合成成功：{output_path}")

    except Exception as e:
        print(f"❌ 合成失败（{left_img_path.name} + {right_img_path.name}）：{str(e)}")


def batch_process_folder(input_folder: str, output_folder: str) -> None:
    """
    批量处理输入文件夹中的图片对，合成后保存到指定输出文件夹
    :param input_folder: 原图片所在的文件夹路径
    :param output_folder: 合成图要保存的目标文件夹路径
    """
    input_dir = Path(input_folder).resolve()
    output_dir = Path(output_folder).resolve()

    if not input_dir.exists() or not input_dir.is_dir():
        print(f"❌ 输入文件夹不存在或无效：{input_dir}")
        return

    # 获取所有带"1.png"后缀的图片（左图候选）
    left_files = sorted(input_dir.glob("*1.png"), key=lambda x: x.name)
    if not left_files:
        print(f"❌ 输入文件夹中未找到带'1.png'的图片：{input_dir}")
        return

    # 遍历左图，匹配对应的右图并合成
    for left_file in left_files:
        # 生成右图文件名（去掉"1"）：例如 "01_Orange1.png" → "01_Orange.png"
        right_filename = left_file.name.replace("1.png", ".png")
        right_file = input_dir / right_filename

        # 检查右图是否存在
        if not right_file.exists():
            print(f"⚠️ 跳过：未找到对应的右图 {right_filename}（左图：{left_file.name}）")
            continue

        # 生成合成图在输出目录中的路径
        combined_filename = left_file.name.replace("1.png", "_combined.png")
        combined_file = output_dir / combined_filename

        # 合成图片
        merge_pair(left_file, right_file, combined_file)

    print("\n📌 批量处理完成！合成图已保存到：", output_dir)


if __name__ == "__main__":
    # 替换为你的输入文件夹和输出文件夹绝对路径
    INPUT_FOLDER = r"../add_text/pomeranian_colors"  # 原图片所在文件夹
    OUTPUT_FOLDER = r"../add_text/pomeranian_colors_combined"  # 合成图要保存的新文件夹
    batch_process_folder(INPUT_FOLDER, OUTPUT_FOLDER)